import sys,os
import numpy as np
import pandas as pd

### Get and set workdirectory -------------------------------------------------
# workdir = 'C://Users//Irrelewandler//Nextcloud//BMI - Burial Mound Inequality//resubmit//addiational_data//plotting//time_series'
workdir = 'D://Cloud//BMI - Burial Mound Inequality//resubmit//addiational_data//plotting//time_series'
os.chdir(workdir)
print ('%s uses working directory %s'%(sys.argv[0],workdir))


# -----------------------------------------------------------------------------
# Load grave and violence data
dfNI = pd.read_csv("data_input/[2023.01.12] graeber_nikulka_mueller_laabs_living.csv")
dfVI = pd.read_csv("data_input/[2023.01.12] gewalt_und_graeber_-8300_1900.csv", sep=";")

# Alter columns to numeric dtype for sorting and querrying
dfNI["ID"] = pd.to_numeric(dfNI["ID"], errors='coerce')
# Sort dfs by "ID"
dfNI = dfNI.sort_values(by="ID")
dfVI = dfVI.sort_values(by="ID")

# Write important violence data into dfNI matching the right "ID"
for i in dfNI.index:
    for j in dfVI.index:
        if dfNI.loc[i,"ID"] == dfVI.loc[j,"ID"]:
            dfNI.loc[i,"n_trauma"] = dfVI.loc[j,"injured_individual_n"]
            dfNI.loc[i,"n_trauma_base"] = dfVI.loc[j,"burial_n"]
        else:
            pass

# Convert specific columns to numertic dtype
lst_num = ["BCE_dating_start","BCE_dating_end","step_200y","burial_n",
           "injured_individual_n"]
for i in lst_num:
    dfVI[i] = pd.to_numeric(dfVI[i], errors='coerce')
# Convert specific columns to numertic dtype
lst_num = ["BCE_CE_start","BCE_CE_end","step_200a","n_burial",
           "ID","n_living"]
for i in lst_num:
    dfNI[i] = pd.to_numeric(dfNI[i], errors='coerce')

# Create df of summing values in 200y steps for time series plotting
df_count = pd.DataFrame(index = \
                        dfNI["step_200a"].value_counts().index.sort_values())
# Calculate sum, mean, median, etc. for wanted values
for i in dfNI["step_200a"].value_counts().index.sort_values():
    mask = np.logical_and(dfNI.loc[::,"step_200a"] == i,
                          dfNI.loc[::,"flat-collect"] == "f")
    df_count.loc[i,"n_flat_sum"] = dfNI.loc[mask,"n_burial"].sum()
    mask = np.logical_and(dfNI.loc[::,"step_200a"] == i,
                          dfNI.loc[::,"flat-collect"] == "c")
    df_count.loc[i,"n_collective_sum"] = dfNI.loc[mask,"n_burial"].sum()
    # df_count.loc[i,"n_burial_mean"] = dfNI.loc[mask,"n_burial"].mean()
    # df_count.loc[i,"n_burial_std"] = dfNI.loc[mask,"n_burial"].std()
    mask3 = np.logical_and(dfNI.loc[::,"step_200a"] == i,
                            dfNI.loc[::,"n_living"].notna())
    df_count.loc[i,"living_sum"] = dfNI.loc[mask3,"n_living"].sum()
    df_count.loc[i,"living_mean"] = dfNI.loc[mask3,"n_living"].mean()
    df_count.loc[i,"living_std"] = dfNI.loc[mask3,"n_living"].std()
# -----------------------------------------------------------------------------


# -----------------------------------------------------------------------------
# Traumata 
for i in dfVI["step_200y"].value_counts().index.sort_values():
    mask = dfVI.loc[::,"step_200y"] == i
    df_count.loc[i,"n_trauma_sum"] = dfVI.loc[mask,"injured_individual_n"].sum()
    df_count.loc[i,"n_trauma_base_sum"] = dfVI.loc[mask,"burial_n"].sum()
# -----------------------------------------------------------------------------


# -----------------------------------------------------------------------------
# Load data of burila mounds
dfBM = pd.read_csv("data_input/burial_mounds_data_200_years.csv")
# Add burial mound data set to count df
for i in dfBM["time_series_dating"].value_counts().index.sort_values():
    mask = dfBM.loc[::,"time_series_dating"] == i
    if i == "3800_3600":
        dfBM.loc[mask,"step_200a"] = -3700
    elif i == "3600_3400":
        dfBM.loc[mask,"step_200a"] = -3500
    elif i == "3400_3200":
        dfBM.loc[mask,"step_200a"] = -3300
    elif i == "3200_3000":
        dfBM.loc[mask,"step_200a"] = -3100
    elif i == "3000_2800":
        dfBM.loc[mask,"step_200a"] = -2900
    elif i == "2800_2600":
        dfBM.loc[mask,"step_200a"] = -2700
    elif i == "2600_2400":
        dfBM.loc[mask,"step_200a"] = -2500
    elif i == "2400_2200":
        dfBM.loc[mask,"step_200a"] = -2300
    elif i == "2200_2000":
        dfBM.loc[mask,"step_200a"] = -2100
    elif i == "2000_1800":
        dfBM.loc[mask,"step_200a"] = -1900
    elif i == "1800_1600":
        dfBM.loc[mask,"step_200a"] = -1700
    elif i == "1600_1400":
        dfBM.loc[mask,"step_200a"] = -1500
    elif i == "1400_1200":
        dfBM.loc[mask,"step_200a"] = -1300
    elif i == "1200_1000":
        dfBM.loc[mask,"step_200a"] = -1100
    elif i == "1000_800":
        dfBM.loc[mask,"step_200a"] = -900
    elif i == "800_600":
        dfBM.loc[mask,"step_200a"] = -700
    elif i == "600_400":
        dfBM.loc[mask,"step_200a"] = -500
    elif i == "400_200":
        dfBM.loc[mask,"step_200a"] = -300
    elif i == "200_0":
        dfBM.loc[mask,"step_200a"] = -100
    else:
        pass
# Add burial mound data set to count df
for i in dfNI["step_200a"].value_counts().index.sort_values():
    mask = dfBM.loc[::,"step_200a"] == i
    # Sum all burial mounds
    df_count.loc[i,"n_mound_sum"] = len(dfBM.loc[mask,::])
    # # Calculate mean size of burial mound cemetery
    # df_count.loc[i,"n_mound_mean"] = \
    #     dfBM.loc[mask,"site_name"].value_counts().mean()
    # df_count.loc[i,"n_mound_std"] = \
    #     dfBM.loc[mask,"site_name"].value_counts().std()
# -----------------------------------------------------------------------------


# -----------------------------------------------------------------------------
# Calculate sums of and ratios between burials & mounds and trauma
for i in dfNI["step_200a"].value_counts().index.sort_values():
    df_count.loc[i,"burial_all_sum"] = \
         (df_count.loc[i,"n_flat_sum"] + df_count.loc[i,"n_mound_sum"] +
          df_count.loc[i,"n_collective_sum"])
    df_count.loc[i,"burial_flat_mount_sum"] = \
         (df_count.loc[i,"n_flat_sum"] + df_count.loc[i,"n_mound_sum"])
    df_count.loc[i,"ratio_flat_all"] = \
         df_count.loc[i,"n_flat_sum"] /\
         (df_count.loc[i,"n_flat_sum"] + df_count.loc[i,"n_mound_sum"] +
          df_count.loc[i,"n_collective_sum"])
    df_count.loc[i,"ratio_mount_all"] = \
         df_count.loc[i,"n_mound_sum"] /\
         (df_count.loc[i,"n_flat_sum"] + df_count.loc[i,"n_mound_sum"] +
          df_count.loc[i,"n_collective_sum"])
    df_count.loc[i,"ratio_flat_all_no_collect"] = \
        df_count.loc[i,"n_flat_sum"] /\
        (df_count.loc[i,"n_flat_sum"] + df_count.loc[i,"n_mound_sum"])
    df_count.loc[i,"ratio_mount_all_no_collect"] = \
        df_count.loc[i,"n_mound_sum"] /\
        (df_count.loc[i,"n_flat_sum"] + df_count.loc[i,"n_mound_sum"])
    df_count.loc[i,"ratio_trauma"] = df_count.loc[i,"n_trauma_sum"] /\
          df_count.loc[i,"n_trauma_base_sum"]     
# -----------------------------------------------------------------------------


# -----------------------------------------------------------------------------
# Load data of burila mounds -> article
dfGI = pd.read_csv("data_input/GINI_volume.csv")
# Add burial mound data set to count df
for i in dfNI["step_200a"].value_counts().index.sort_values():
    for j in dfGI.index:
        if i == dfGI.loc[j,"BCE"]:
            df_count.loc[i,"Gini"] = dfGI.loc[j,"GINI"]
            df_count.loc[i,"V_mean"] = dfGI.loc[j,"V_mean"]
            df_count.loc[i,"V_std"] = dfGI.loc[j,"V_std"]
            df_count.loc[i,"V_min"] = dfGI.loc[j,"V_min"]
            df_count.loc[i,"V_max"] = dfGI.loc[j,"V_max"]
        else:
            pass
# -----------------------------------------------------------------------------        


# -----------------------------------------------------------------------------
# Write df to CSV
df_count.loc[-4000:0,::].to_csv("data_output/df_200a_-4000_0.csv")       
# -----------------------------------------------------------------------------
     